# Copyright 2010 Thomas Conneely
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import scipy.integrate as integrate
import scipy as spy
from math import sqrt

import pdb

class Polynomial(object):
    '''Polynomial class, takes a function that does some
    maths on a value x and returns the answer, e.g.:

        def fn(x):
            return x**2
    
    Polynomial acts like a wrapper for fn(x) so you can add useful
    functions'''

    def __init__(self, func):
        self.func = func
        self.normalised = False

    def __call__(self, x):
        return float(self.func(x))

    def _squared(self, x):
        return float(self.func(x) ** 2)


class NormalisedPoly(Polynomial):
    def __init__(self, func, min=0, max=100):
        Polynomial.__init__(self, func)
        self.normalised = True
        self.min = min
        self.max = max

        # First normalise the polynomial
        poly_integral = integrate.quad(self.func, min, max)
        # Now normalisation co-eff for poly (A) is given by
        self.norm_const = 1 / poly_integral[0]

    
    def __call__(self, x):
        return self.norm_const * self.func(x)


class Bin(object):
    def __init__(self, poly, prob, min, max, prob_min, prob_max):
        self.prob = prob
        # use non-normalised polynomial for calculating values
        self.poly = poly.func
        
        # min is the minimum x-value of the bin
        # similarly for max
        self.min = min
        self.max = max

        # These variables are the bins start and
        # end in 'probability space'
        self.prob_min = prob_min
        self.prob_max = prob_max

    def inrange(self, x):
        '''Returns True if a random number is in
        the range corresponding to this bin'''

        if (x >= self.prob_min) and (x <= self.prob_max):
            return True
        else:
            return False
    
    def rand(self):
        '''Generate a random number within the 
        x-range of the bin (uniform distribution)'''

        return spy.random.uniform(self.min, self.max)

class poly_random(object):
    '''A random number generator, which produces numbers according 
    to a specified distribution'''
    def __init__(self, poly, min=0, max=100, no_bins=1000):
        self.poly = NormalisedPoly(poly, min, max)
        
        self.bins = []
        sum = 0

        bin_width = (max - min) / no_bins
        for m in spy.arange(min, max, bin_width):
            bin_min = m
            bin_max = (m + bin_width) 
            prob = integrate.quad(self.poly, bin_min, bin_max)[0]
    
            prob_min = sum
            prob_max = sum + prob

            sum += prob

            #print prob_min, prob_max, sum

            self.bins.append(Bin(self.poly, prob, bin_min, bin_max, prob_min, prob_max))

    def __call__(self):
        r = spy.random.uniform()
        for b in self.bins:
            if b.inrange(r):
                return b.rand()
               
    def _ValueInBin(self, bin):
        '''Tests whether a value is in range of a specific bin'''
        if (x >= bin.prob_min) and (x < bin.prob_max):
            return True
        else:
            return False
    
        

def test_poly(x):
    import math
    return (math.sin(x)) ** 6
    
if __name__=='__main__':
    import pylab
    import math

    p = []
    r = poly_random(test_poly, min=0, max=math.pi, no_bins=100)
    
    for i in range(0,100000):
        p.append(r())
    
    histo = spy.histogram(p, 100)
    pylab.plot(histo[1], histo[0])
    pylab.show()

